{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "import mxnet as mx\n",
    "from mxnet.gluon import nn\n",
    "from mxnet import gluon, nd, autograd\n",
    "import logging\n",
    "import os\n",
    "import glob\n",
    "from PIL import Image\n",
    "%matplotlib inline\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from datetime import timedelta"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tqdm import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "import yaml\n",
    "from os.path import join\n",
    "from time import sleep"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "sys.path = ['/home/hzh/.local/lib/python3.5/site-packages/'] + sys.path\n",
    "import cv2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create logger with 'spam_application'\n",
    "logger = logging.getLogger('VGG_net')\n",
    "logger.setLevel(logging.DEBUG)\n",
    "# create file handler which logs even debug messages\n",
    "fh = logging.FileHandler('classify.log')\n",
    "fh.setLevel(logging.DEBUG)\n",
    "fh.setFormatter(logging.Formatter('%(asctime)s - %(message)s'))\n",
    "ch = logging.StreamHandler()\n",
    "\n",
    "logger.addHandler(fh)\n",
    "logger.addHandler(ch)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_classes = 3753\n",
    "ctx = mx.cpu()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_net(num_outputs):\n",
    "    def vgg_block(num_convs, channels):\n",
    "        out = nn.Sequential()\n",
    "        for _ in range(num_convs):\n",
    "            out.add(nn.Conv2D(channels=channels, kernel_size=3,\n",
    "                          padding=1, activation='relu'))\n",
    "        out.add(nn.MaxPool2D(pool_size=2, strides=2))\n",
    "        return out\n",
    "\n",
    "    def vgg_stack(architecture):\n",
    "        out = nn.Sequential()\n",
    "        for (num_convs, channels) in architecture:\n",
    "            out.add(vgg_block(num_convs, channels))\n",
    "        return out\n",
    "    architecture = ((1,64), (1,128), (2,256), (2,512))\n",
    "    net = nn.Sequential()\n",
    "    with net.name_scope():\n",
    "        net.add(vgg_stack(architecture))\n",
    "        net.add(nn.Flatten())\n",
    "        net.add(nn.Dense(512, activation=\"relu\"))\n",
    "        net.add(nn.Dropout(.5))\n",
    "        net.add(nn.Dense(512, activation=\"relu\"))\n",
    "        net.add(nn.Dropout(.5))\n",
    "        net.add(nn.Dense(num_outputs))\n",
    "    return net"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "net2 = create_net(num_classes)\n",
    "net2.load_parameters(\"./net.params\", ctx=ctx)    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_db(fp):\n",
    "    with open(fp, 'r') as f:\n",
    "        return yaml.load(f, Loader=yaml.FullLoader)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "db_index = load_db('./chi_seq_imgs/index.yaml')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(\"./data/chi3500.txt\", 'r') as f:\n",
    "    chi3500 = f.readline()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1587/1587 [02:36<00:00,  3.52it/s]\n"
     ]
    }
   ],
   "source": [
    "cnt = 0\n",
    "sbt_seq = []\n",
    "for k in tqdm(db_index.keys()):\n",
    "    info = db_index[k]\n",
    "    chi_chain = []\n",
    "    for fn in info['chi_seq']:\n",
    "        img = cv2.imread(join('./chi_seq_imgs', fn))\n",
    "        img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)\n",
    "        img = cv2.resize(img, (50, 50))\n",
    "        img_nd =  mx.nd.array(img)\n",
    "#         print(\"img shape=\", img_nd.shape)\n",
    "        b = nd.reshape(img_nd, (1, 1, 50, 50))\n",
    "        c = b.astype(np.float32)\n",
    "        output = net2(c.as_in_context(ctx))\n",
    "        predictions = nd.argmax(output, axis=1)\n",
    "#         print(\"prediction=\", predictions)\n",
    "#         plt.figure()\n",
    "#         plt.imshow(img)\n",
    "        c = chi3500[int(predictions.asnumpy()[0])]\n",
    "        chi_chain.append(c)\n",
    "    sbt_seq.append((info['time'], ''.join(chi_chain)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "outputs": [],
   "source": [
    "def time2str(t):\n",
    "    dt = timedelta(seconds=t/1000.0)\n",
    "    s = str(dt)\n",
    "    return s.replace('.', ',')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "metadata": {},
   "outputs": [],
   "source": [
    "def dump2srt(seq, fp=\"./a.srt\"):\n",
    "    lines = []\n",
    "    for i in tqdm(range(len(seq))):\n",
    "        lines.append(\"%d\"%i)\n",
    "        cur_tm = seq[i][0]\n",
    "        next_tm = seq[i+1][0] if i+1 < len(seq) else cur_tm+1000 # plus 1 seconds\n",
    "        lines.append(\"%s --> %s\" % (time2str(cur_tm), time2str(next_tm)) )\n",
    "        lines.append(\"%s\" % seq[i][1])\n",
    "    with open(fp, 'w') as f:\n",
    "        f.write('\\n'.join(lines))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1587/1587 [00:00<00:00, 81188.99it/s]\n"
     ]
    }
   ],
   "source": [
    "dump2srt(sbt_seq)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pytorch_venv",
   "language": "python",
   "name": "pytorch_venv"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
